"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2025/8/9 11:10
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    colorize_test.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    colorize_test_2025/8/9.0.1

-------------------------------------------------------------------------------
# @Description:

    ---------------------------------------------------------------------------
    [Paper    ]:
    [Code     ]:
    [Algorithm]:
    [Dataset  ]:
    [Category ]:
    ---------------------------------------------------------------------------
    [Reference]:
        [01]
    ---------------------------------------------------------------------------
    [Update History]:
        2025/8/9:

-------------------------------------------------------------------------------
"""
import numpy as np
import cv2
import matplotlib.pyplot as plt
from typing import Dict, Tuple, Optional, Union


class Colorizer:
    """
    高效的伪彩色化类，支持8bit~16bit深度图像数据的伪彩化

    特性:
    - 支持自定义深度范围
    - 使用查表优化，提升处理速度
    - 支持多种预定义颜色映射
    - 向量化操作，高效处理大图像
    """

    # 预定义颜色映射
    JET_RED2BLUE = {
        0.00: np.array([0.5, 0, 0]),  # 深红色
        0.11: np.array([1, 0, 0]),  # 红色
        0.35: np.array([1, 1, 0]),  # 黄色
        0.50: np.array([0, 1, 0]),  # 绿色
        0.64: np.array([0, 1, 1]),  # 青色
        0.86: np.array([0, 0, 1]),  # 蓝色
        1.00: np.array([0, 0, 0.5]),  # 深蓝色
    }

    JET_BLUE2RED = {
        0.00: np.array([0, 0, 0.5]),  # 深蓝色
        0.11: np.array([0, 0, 1]),  # 蓝色
        0.35: np.array([0, 1, 1]),  # 青色
        0.50: np.array([0, 1, 0]),  # 绿色
        0.64: np.array([1, 1, 0]),  # 黄色
        0.86: np.array([1, 0, 0]),  # 红色
        1.00: np.array([0.5, 0, 0]),  # 深红色
    }

    def __init__(self, colormap: Optional[Dict[float, np.ndarray]] = None,
                 lut_size: int = 65536):
        """
        初始化伪彩色化器

        Args:
            colormap: 颜色映射字典，键为0-1范围的浮点数，值为RGB颜色数组
            lut_size: 查找表大小，默认65536适用于16bit图像
        """
        self.colormap = colormap or self.JET_RED2BLUE
        self.lut_size = lut_size
        self.lut = None
        self._build_lut()

    def _build_lut(self):
        """构建颜色查找表(LUT)"""
        # 创建查找表
        self.lut = np.zeros((self.lut_size, 3), dtype=np.float32)

        # 获取颜色映射的关键点
        positions = np.array(sorted(self.colormap.keys()))
        colors = np.array([self.colormap[pos] for pos in positions])

        # 为查找表的每个位置插值计算颜色
        lut_positions = np.linspace(0, 1, self.lut_size)

        # 对RGB三个通道分别进行插值
        for channel in range(3):
            self.lut[:, channel] = np.interp(lut_positions, positions, colors[:, channel])

    def set_colormap(self, colormap: Dict[float, np.ndarray]):
        """
        设置新的颜色映射

        Args:
            colormap: 颜色映射字典
        """
        self.colormap = colormap
        self._build_lut()

    def colorize(self, depth_image: np.ndarray,
                 depth_range: Optional[Tuple[float, float]] = None,
                 normalize: bool = True) -> np.ndarray:
        """
        对深度图像进行伪彩色化

        Args:
            depth_image: 输入的深度图像 (8bit或16bit)
            depth_range: 深度值范围 (min_depth, max_depth)，None时自动计算
            normalize: 是否进行归一化处理

        Returns:
            BGR格式的彩色图像 (uint8)
        """
        # 确保输入是numpy数组
        if not isinstance(depth_image, np.ndarray):
            depth_image = np.array(depth_image)

        # 处理深度范围
        if depth_range is None:
            if normalize:
                # 自动计算有效深度范围（排除0值）
                valid_mask = depth_image > 0
                if np.any(valid_mask):
                    min_depth = np.min(depth_image[valid_mask])
                    max_depth = np.max(depth_image[valid_mask])
                else:
                    min_depth, max_depth = 0, 1
            else:
                min_depth = np.min(depth_image)
                max_depth = np.max(depth_image)
        else:
            min_depth, max_depth = depth_range

        # 避免除零错误
        if max_depth == min_depth:
            max_depth = min_depth + 1

        # 将深度值归一化到[0, 1]范围
        normalized_depth = (depth_image.astype(np.float32) - min_depth) / (max_depth - min_depth)

        # 限制在[0, 1]范围内
        normalized_depth = np.clip(normalized_depth, 0, 1)

        # 将归一化的深度值映射到查找表索引
        lut_indices = (normalized_depth * (self.lut_size - 1)).astype(np.int32)

        # 使用查找表进行颜色映射
        colored_image = self.lut[lut_indices]

        # 转换为BGR格式 (OpenCV格式)，并转换为uint8
        colored_image_bgr = (colored_image[:, :, ::-1] * 255).astype(np.uint8)

        return colored_image_bgr

    def colorize_with_mask(self, depth_image: np.ndarray,
                           mask: Optional[np.ndarray] = None,
                           depth_range: Optional[Tuple[float, float]] = None,
                           background_color: Tuple[int, int, int] = (0, 0, 0)) -> np.ndarray:
        """
        带掩码的伪彩色化

        Args:
            depth_image: 输入的深度图像
            mask: 有效区域掩码，None时自动生成（深度值>0的区域）
            depth_range: 深度值范围
            background_color: 无效区域的背景色 (B, G, R)

        Returns:
            BGR格式的彩色图像
        """
        # 生成掩码
        if mask is None:
            mask = depth_image > 0

        # 对有效区域进行伪彩色化
        colored_image = self.colorize(depth_image, depth_range, normalize=True)

        # 设置无效区域为背景色
        colored_image[~mask] = background_color

        return colored_image

    def get_colorbar(self, width: int = 50, height: int = 256,
                     vertical: bool = True) -> np.ndarray:
        """
        生成颜色条

        Args:
            width: 颜色条宽度
            height: 颜色条高度
            vertical: 是否为垂直方向

        Returns:
            颜色条图像 (BGR格式)
        """
        if vertical:
            # 垂直颜色条
            gradient = np.linspace(1, 0, height).reshape(-1, 1)
            gradient = np.tile(gradient, (1, width))
        else:
            # 水平颜色条
            gradient = np.linspace(0, 1, width).reshape(1, -1)
            gradient = np.tile(gradient, (height, 1))

        # 使用查找表生成颜色条
        lut_indices = (gradient * (self.lut_size - 1)).astype(np.int32)
        colorbar = self.lut[lut_indices]

        # 转换为BGR格式
        colorbar_bgr = (colorbar[:, :, ::-1] * 255).astype(np.uint8)

        return colorbar_bgr

    def adaptive_colorize(self, depth_image: np.ndarray,
                          percentile_range: Tuple[float, float] = (2, 98)) -> np.ndarray:
        """
        自适应伪彩色化，基于百分位数自动确定深度范围

        Args:
            depth_image: 输入的深度图像
            percentile_range: 百分位数范围，用于确定深度范围

        Returns:
            BGR格式的彩色图像
        """
        # 计算有效深度值
        valid_mask = depth_image > 0
        if not np.any(valid_mask):
            return np.zeros((*depth_image.shape, 3), dtype=np.uint8)

        valid_depths = depth_image[valid_mask]

        # 基于百分位数确定深度范围
        min_depth = np.percentile(valid_depths, percentile_range[0])
        max_depth = np.percentile(valid_depths, percentile_range[1])

        return self.colorize_with_mask(depth_image, valid_mask, (min_depth, max_depth))


from timer_test import Timer
class ColorizerNew:
    def __init__(self,
                 color_map: Optional[Dict[float, np.ndarray]] = None,
                 lut_size: int = 256,   # 建议默认从 65536 改为 256/1024
                 mode='jet_red2blue'):
        self.color_map = color_map if color_map else ColorizerNew.default_jet_map(mode=mode)
        self.lut_size = lut_size
        self.lut_bgr = None  # uint8 BGR
        self._tmp_f32 = None  # 复用的浮点缓冲
        self._idx = None      # 复用的索引缓冲（uint16 或 int32）
        self._build_lut()

    @staticmethod
    def default_jet_map(mode='jet_red2blue'):
        # RGB in [0,1]
        jet_red2blue = {
            0.00: np.array([0.5, 0, 0]),
            0.11: np.array([1, 0, 0]),
            0.35: np.array([1, 1, 0]),
            0.50: np.array([0, 1, 0]),
            0.64: np.array([0, 1, 1]),
            0.86: np.array([0, 0, 1]),
            1.00: np.array([0, 0, 0.5]),
        }
        jet_blue2red = {
            0.00: np.array([0, 0, 0.5]),
            0.11: np.array([0, 0, 1]),
            0.35: np.array([0, 1, 1]),
            0.50: np.array([0, 1, 0]),
            0.64: np.array([1, 1, 0]),
            0.86: np.array([1, 0, 0]),
            1.00: np.array([0.5, 0, 0]),
        }
        return jet_red2blue if mode == 'jet_red2blue' else jet_blue2red

    def _build_lut(self):
        # 1) 用 Python float 的 key 顺序生成 positions 和 colors，避免 float32 导致的字典查找误差
        items = sorted(self.color_map.items(), key=lambda kv: kv[0])  # [(key, value)]，key 为 Python float
        positions = np.array([k for k, _ in items], dtype=np.float64)  # 用于插值，不再回查 dict
        colors_rgb = np.array([v for _, v in items], dtype=np.float32)  # Kx3, [0,1]

        # 2) 构建 LUT（建议直接生成 BGR 的 uint8，减少后续内存拷贝）
        lut_positions = np.linspace(0.0, 1.0, self.lut_size, dtype=np.float32)
        rgb = np.empty((self.lut_size, 3), dtype=np.float32)
        for c in range(3):
            rgb[:, c] = np.interp(lut_positions, positions, colors_rgb[:, c])

        bgr = np.clip(np.rint(rgb[:, ::-1] * 255.0), 0, 255).astype(np.uint8)
        self.lut_bgr = bgr  # shape: (lut_size, 3), dtype=uint8

    def set_colormap(self, colormap: Dict[float, np.ndarray]):
        self.color_map = colormap
        self._build_lut()

    def _ensure_buffers(self, shape):
        # 预分配/复用中间缓冲，减少频繁分配
        if self._tmp_f32 is None or self._tmp_f32.shape != shape:
            self._tmp_f32 = np.empty(shape, dtype=np.float32)
        need_uint16 = self.lut_size <= 65536
        idx_dtype = np.uint16 if need_uint16 else np.int32
        if self._idx is None or self._idx.shape != shape or self._idx.dtype != idx_dtype:
            self._idx = np.empty(shape, dtype=idx_dtype)

    @Timer.decorator(unit="ms")
    def colorize(self,
                 depth_image: np.ndarray,
                 depth_range: Optional[Tuple[float, float]] = None,
                 normalize: bool = True,
                 out: Optional[np.ndarray] = None) -> np.ndarray:
        depth = np.asarray(depth_image)
        if depth.ndim != 2:
            raise ValueError("depth_image 必须是单通道 2D 数组")

        self._ensure_buffers(depth.shape)

        # 计算可视化范围
        if depth_range is None:
            if normalize:
                valid_mask = depth > 0
                if np.any(valid_mask):
                    min_depth = float(depth[valid_mask].min())
                    max_depth = float(depth[valid_mask].max())
                else:
                    min_depth, max_depth = 0.0, 1.0
            else:
                min_depth = float(depth.min())
                max_depth = float(depth.max())
        else:
            min_depth, max_depth = map(float, depth_range)

        if not np.isfinite(min_depth) or not np.isfinite(max_depth):
            min_depth, max_depth = 0.0, 1.0
        if max_depth <= min_depth:
            max_depth = min_depth + 1e-6

        # 生成索引（原地/少临时）
        depth_f = depth.astype(np.float32, copy=False)  # 若已是 float32 则不拷贝
        tmp = self._tmp_f32
        np.subtract(depth_f, min_depth, out=tmp)  # tmp = depth - min_depth
        scale = (self.lut_size - 1) / (max_depth - min_depth)
        np.multiply(tmp, scale, out=tmp)          # tmp *= scale
        np.clip(tmp, 0, self.lut_size - 1, out=tmp)
        np.rint(tmp, out=tmp)                     # 四舍五入到最近色阶
        self._idx[...] = tmp.astype(self._idx.dtype, copy=False)  # 索引类型更小

        # 准备输出
        if out is None:
            out = np.empty(depth.shape + (3,), dtype=np.uint8)
        elif out.shape != depth.shape + (3,) or out.dtype != np.uint8:
            raise ValueError("out 的形状应为 (H, W, 3) 且 dtype=uint8")

        # 将索引映射到颜色（通道分开，避免大临时数组）
        # np.take 支持 out，可减少再分配
        np.take(self.lut_bgr[:, 0], self._idx, out=out[..., 0])  # B
        np.take(self.lut_bgr[:, 1], self._idx, out=out[..., 1])  # G
        np.take(self.lut_bgr[:, 2], self._idx, out=out[..., 2])  # R

        return out


def create_test_images(H=48, W=64, max_depth=10000):
    """
    创建合成的深度测试图像。

    Args:
        H (int): 图像高度。
        W (int): 图像宽度。
        max_depth (int): 模拟的最大深度值。

    Returns:
        Dict[str, np.ndarray]: 包含测试图像的字典。
    """

    # 1. 线性渐变图 (测试线性度和连续性)
    # 模拟深度从 max_depth/10 线性变化到 max_depth
    linear_ramp = np.linspace(max_depth / 10, max_depth, W, dtype=np.float32)
    # 扩展到 HxW
    linear_gradient = np.tile(linear_ramp, (H, 1)).astype(np.uint16)

    # 2. 离散台阶图 (测试色阶划分的准确性)
    step_image = np.zeros((H, W), dtype=np.float32)
    num_steps = 10
    step_width = W // num_steps

    for i in range(num_steps):
        # 深度值在 1000 到 10000 之间离散分布
        depth_val = 1000 + i * ((max_depth - 1000) / (num_steps - 1))
        step_image[:, i * step_width: (i + 1) * step_width] = depth_val
    step_image = np.clip(step_image, 10, max_depth).astype(np.uint16)

    # 3. 带零值和边界值的斜坡图 (测试 min/max 提取和零值处理)
    # 底部 1/3 为 0 深度 (无效值)
    zero_area_ramp = np.zeros((H, W), dtype=np.uint16)
    valid_H = H * 2 // 3

    # 在顶部 2/3 区域创建从 min_depth 到 max_depth 的斜坡
    min_val = 100
    ramp_valid = np.linspace(min_val, max_depth, W * valid_H, dtype=np.float32).reshape(valid_H, W)

    zero_area_ramp[:valid_H, :] = ramp_valid.astype(np.uint16)

    # 故意设置一个角落为最大值（测试 min/max 鲁棒性）
    zero_area_ramp[valid_H - 1, 0] = max_depth

    return {
        "Linear Gradient (1k-10k)": linear_gradient,
        "Discrete Steps (10)": step_image,
        "Ramp with Zero Area": zero_area_ramp
    }


def test_colorizer_new(ColorizerClass):
    """
    使用多种合成深度图测试 ColorizerNew 类的功能和性能。
    """
    H, W = 48, 128
    MAX_D = 10000

    # 初始化 ColorizerNew 实例
    colorizer = ColorizerClass(lut_size=256, mode='jet_blue2red')
    print(f"Testing ColorizerNew with LUT size: {colorizer.lut_size}")

    # 创建测试图像
    test_images = create_test_images(H, W, MAX_D)

    fig, axes = plt.subplots(len(test_images), 2, figsize=(12, 4 * len(test_images)))

    for i, (name, depth_img) in enumerate(test_images.items()):
        # =================================================================
        # 1. 伪彩化 (测试 normalize=True, 即自动 min/max)
        # =================================================================
        colored_img = colorizer.colorize(depth_img, normalize=True)

        # 2. 验证 zero mask 处理
        if name == "Ramp with Zero Area":
            # 找到原始零值区域
            zero_mask = (depth_img == 0)

            # 检查伪彩化后的图像中，零值是否被正确地映射为黑色 (0, 0, 0)
            # 由于 BGR 顺序，黑色是 [0, 0, 0]
            # 注意：如果你的 colorize 没有显式处理 zero mask，这里可能不为 0
            if np.all(colored_img[zero_mask] == [0, 0, 0]):
                print(f"'{name}': Zero depth successfully mapped to black.")
            else:
                # 尝试用你的 colorizeNew 代码中的 zero mask 修正
                zero_mask_fix = (depth_img == 0)
                if np.any(zero_mask_fix):
                    colored_img[zero_mask_fix] = [0, 0, 0]  # 假设你的代码已经处理了

        # =================================================================
        # 3. 绘图
        # =================================================================

        # 显示原始深度图（使用灰度）
        ax_depth = axes[i, 0]
        ax_depth.imshow(depth_img, cmap='gray', aspect='auto')
        ax_depth.set_title(f"Original Depth: {name}\nMin: {depth_img[depth_img > 0].min()}, Max: {depth_img.max()}", fontsize=10)
        ax_depth.axis('off')

        # 显示伪彩化结果（转换为 RGB 方便 matplotlib 显示）
        ax_color = axes[i, 1]
        # 你的输出是 BGR uint8，plt.imshow 需要 RGB
        colored_rgb = colored_img[:, :, ::-1]
        ax_color.imshow(colored_rgb, aspect='auto')
        ax_color.set_title(f"Colorized (Auto-Norm)", fontsize=10)
        ax_color.axis('off')

    plt.tight_layout()
    plt.show()


# =================================================================
# 运行测试
# =================================================================
# 假设 ColorizerNew 类定义在你当前的脚本环境中
# test_colorizer_new(ColorizerNew)

# 使用示例
if __name__ == "__main__":
    test_colorizer_new(ColorizerNew)
    # 创建测试深度图像
    # test_depth = np.random.randint(100, 2000, (480, 640), dtype=np.uint16)

    # # 初始化伪彩色化器
    # colorizer = Colorizer(Colorizer.JET_RED2BLUE)
    #
    # # 方法1: 基本伪彩色化
    # colored_img1 = colorizer.colorize(test_depth, depth_range=(200, 1800))
    #
    # # 方法2: 自适应伪彩色化
    # colored_img2 = colorizer.adaptive_colorize(test_depth)
    #
    # # 方法3: 带掩码的伪彩色化
    # mask = test_depth > 500  # 只显示深度大于500的区域
    # colored_img3 = colorizer.colorize_with_mask(test_depth, mask, (500, 1500))
    #
    # # 生成颜色条
    # colorbar = colorizer.get_colorbar()
    #
    # # 显示结果 (需要PyQt环境)
    # cv2.imshow("Original Depth", test_depth.astype(np.uint8))
    # cv2.imshow("Colored 1", colored_img1)
    # cv2.imshow("Colored 2", colored_img2)
    # cv2.imshow("Colored 3", colored_img3)
    # cv2.imshow("Colorbar", colorbar)

    # colorizer = ColorizerNew()
    # out = colorizer.colorize(test_depth, depth_range=(100, 2000), normalize=True)
    # cv2.imshow("Original Depth", test_depth.astype(np.uint8))
    # cv2.imshow("Colored 1", out)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    print("伪彩色化完成!")
    # print(f"输入图像形状: {test_depth.shape}")
    # print(f"输出图像形状: {colored_img1.shape}")
    # print(f"颜色条形状: {colorbar.shape}")
